package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"log"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"
	"unicode"

	"gitee.com/shupkg/bolt"
	"github.com/abiosoft/ishell"
	"github.com/spf13/pflag"
)

func main() {
	log.SetFlags(0)
	if len(os.Args) != 2 && len(os.Args) != 3 {
		log.Fatal("use: boltsh <boltdb file>")
	}

	s, err := newBolt(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	defer s.Close()

	if len(os.Args) == 3 {
		switch os.Args[2] {
		case "info":
			fmt.Println(printStruct(s.bolt.DB().Stats()))
		}
		return
	}

	shell := ishell.New()
	shell.SetPrompt("$:/> ")
	shell.AutoHelp(true)
	shell.Interrupt(func(c *ishell.Context, count int, _ string) {
		switch count {
		case 4:
			c.Println("press again to exit")
		case 5:
			c.Stop()
		}
	})

	s.addChdir(shell)
	s.addGet(shell)
	s.addIncr(shell)
	s.addInfo(shell)
	s.addList(shell)
	s.addMkdir(shell)
	s.addPut(shell)
	s.addPwd(shell)
	s.addRm(shell)
	s.addSequence(shell)
	s.addSetSeq(shell)

	shell.ClearScreen()
	shell.Println("BoltDB Bolt: https://gitee.com/shupkg/boltsh")
	shell.Println("----------------------------------------------")
	shell.Run()
}

func newBolt(fn string) (*Bolt, error) {
	db, err := bolt.OpenDefault(fn)
	if err != nil {
		return nil, err
	}
	s := &Bolt{bolt: db}
	return s, nil
}

type Bolt struct {
	bolt    *bolt.DB
	current [][]byte
}

func (s *Bolt) Close() error {
	return s.bolt.Close()
}

func (s *Bolt) openCurrent(tx bolt.Tx) (bolt.Tx, error) {
	return openBucket(tx, s.current)
}

func (s *Bolt) addInfo(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:      "info",
		Help:      "display database info",
		Completer: s.getKeyCompleter(nil),
		Func: wrapCmd(func(c *ishell.Context) error {
			if len(c.Args) == 0 {
				return s.bolt.View(func(tx bolt.Tx) (err error) {
					if tx, err = s.openCurrent(tx); err != nil {
						return
					}
					c.Println("Current Info:")
					c.Println("--------------------")
					c.Println(printStruct(tx.Stats()))
					return nil
				})
			}
			if len(c.Args) == 1 {
				switch c.Args[0] {
				case "tx":
					return s.bolt.View(func(tx bolt.Tx) error {
						c.Println("Tx Info:")
						c.Println("--------------------")
						c.Println(printStruct(tx.Root().Stats()))
						return nil
					})
				case "db":
					c.Println("DB Info:")
					c.Println("--------------------")
					c.Println(printStruct(s.bolt.DB().Stats()))
					return nil
				}
			}
			return fmt.Errorf("usage: info [db | tx]")
		}),
	})
}

func (s *Bolt) addList(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:      "ls",
		Aliases:   []string{"l", "la"},
		Help:      "list buckets or keys under current bucket",
		Completer: s.getKeyCompleter(nil),
		Func:      wrapCmd(s.list),
	})
}

func (s *Bolt) list(c *ishell.Context) error {
	var listOptions []bolt.ListOption
	var (
		bucketOnly bool
		keyOnly    bool
		reverse    bool
		skip       int
		limit       = 100
	)

	if len(c.Args) > 0 {
		pflag.ErrHelp = errors.New("")
		set := pflag.NewFlagSet("list", pflag.ContinueOnError)
		set.BoolVarP(&bucketOnly, "bucket-only", "b", bucketOnly, "")
		set.BoolVarP(&keyOnly, "key-only", "k", keyOnly, "")
		set.BoolVarP(&reverse, "reverse", "r", reverse, "")
		set.IntVarP(&skip, "skip", "o", skip, "")
		set.IntVarP(&limit, "limit", "l", limit, "")
		if err := set.Parse(c.Args); err != nil {
			return nil
		}
		if len(set.Args()) > 0 {
			listOptions = append(listOptions, bolt.ListSeek(toB(set.Arg(0))))
		}
		if bucketOnly {
			listOptions = append(listOptions, bolt.ListOnlyBucket())
		}
		if keyOnly {
			listOptions = append(listOptions, bolt.ListNoBucket())
		}
		if reverse {
			listOptions = append(listOptions, bolt.ListReverse())
		}

		//for i := 0; i < len(c.Args); i++ {
		//	switch c.Args[i] {
		//	case "-b", "--bucket-only":
		//		listOptions = append(listOptions, bolt.ListOnlyBucket())
		//	case "-k", "--key-only":
		//		listOptions = append(listOptions, bolt.ListNoBucket())
		//	default:
		//		listOptions = append(listOptions, bolt.ListSeek(toB(c.Args[i])))
		//	}
		//}
	}

	c.Println("--------------------")
	set, err := s.lists(listOptions...)
	if len(set) > 0 {
		max := 0
		for i, e := range set {
			if i < skip {
				continue
			}
			if i >= skip+limit {
				break
			}
			if max < len(e.Key) {
				max = len(e.Key)
			}
		}
		//if  skip > 0 {
		//
		//}
		//if limit > 0 {
		//
		//}
		sMax := strconv.Itoa(max + 2)
		for i, e := range set {
			if i < skip {
				continue
			}
			if i >= skip+limit {
				break
			}
			if e.Bucket != nil {
				c.Printf("%-"+sMax+"s = %s\n", toS(e.Key), "[BUCKET]")
			} else {
				v := toS(e.Value)
				if len(v) > 50 {
					v = v[:50] + "..."
				}
				c.Printf("%-5d %-"+sMax+"s = %s\n", i+1, toS(e.Key), v)
			}
		}
	}
	return err
}

func (s *Bolt) addChdir(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:      "cd",
		Aliases:   []string{"chdir"},
		Help:      "change to another bucket",
		Completer: s.getKeyCompleter([]string{".."}, bolt.ListOnlyBucket()),
		Func:      wrapCmd(s.chdir, false),
	})
}

func (s *Bolt) chdir(c *ishell.Context) error {
	if len(c.Args) != 1 {
		return fmt.Errorf("usage: cd or chdir <bucket>")
	}
	path := joinPath(s.current, c.Args[0])
	if len(path) == 0 {
		s.current = nil
		c.SetPrompt("$:/> ")
		return nil
	}
	return s.bolt.View(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		switch c.Args[0] {
		case ".":
		case "..":
			tx = tx.Parent()
		default:
			tx = tx.Bucket(toB(c.Args[0]))
		}
		if tx == nil {
			return notFound(path)
		}
		s.current = path
		c.SetPrompt("$:" + sprintPath(path) + "/> ")
		return nil
	})
}

func (s *Bolt) addMkdir(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:    "mkdir",
		Aliases: []string{"create"},
		Help:    "mkdir <bucket> or create <bucket>: create a sub bucket",
		Func:    wrapCmd(s.mkdir, true),
	})
}

func (s *Bolt) mkdir(c *ishell.Context) error {
	if len(c.Args) == 0 {
		return fmt.Errorf("usage: mkdir or create <...bucket>")
	}
	return s.bolt.Update(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		for _, path := range c.Args {
			c.Println("create bucket:", path)
			if _, err = tx.CreateBucket(toB(path)); err != nil {
				break
			}
		}
		return
	})
}

func (s *Bolt) addPut(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:    "put",
		Aliases: []string{"set"},
		Help:    "put key value",
		Func:    wrapCmd(s.put, true),
	})
}

func (s *Bolt) put(c *ishell.Context) error {
	if len(c.Args) < 2 {
		return fmt.Errorf("usage: set or put <key1> <value1> ... <keyN> <valueN>")
	}

	return s.bolt.Update(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		for i := 0; i+1 < len(c.Args); i += 2 {
			c.Printf("put: %-10s -> %s\n", c.Args[i], c.Args[i+1])
			if err = tx.Put(toB(c.Args[i]), toB(c.Args[i+1])); err != nil {
				break
			}
		}
		return
	})
}

func (s *Bolt) addGet(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:      "get",
		Help:      "get [string|time|int] <key>: show time,int or string value by key",
		Func:      wrapCmd(s.get, false),
		Completer: s.getKeyCompleter([]string{"string", "int", "time"}),
	})
}

func (s *Bolt) get(c *ishell.Context) error {
	if len(c.Args) != 1 && len(c.Args) != 2 {
		return fmt.Errorf("usage: get [i[nt] | s[tring] | t[ime]] <key>")
	}

	var (
		key = c.Args[len(c.Args)-1]
		tp  = ""
	)

	if len(c.Args) == 2 {
		switch c.Args[0] {
		case "i", "int", "s", "string", "t", "time":
			tp = c.Args[0][:1]
		default:
			return fmt.Errorf("usage: get [i[nt] | s[tring] | f[loat] | b[ool]] <key1>")
		}
	}

	return s.bolt.View(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		value := tx.Get(toB(key))
		if len(value) == 0 {
			c.Println("<nil>")
			return
		}

		switch tp {
		case "":
			c.Println(toS(value))
		case "s":
			c.Println(string(value))
		case "i":
			id, _ := binary.Uvarint(value)
			c.Printf("%d\n", id)
		case "t":
			var t time.Time
			if err = t.UnmarshalBinary(value); err == nil {
				c.Println(t.Format(time.RFC3339))
			}
		}
		return
	})
}

func (s *Bolt) addRm(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:      "rm",
		Aliases:   []string{"del"},
		Help:      "rm -f <bucket>, rm <key>: delete a bucket or key",
		Completer: s.getKeyCompleter(nil),
		Func:      wrapCmd(s.rm, true),
	})
}

func (s *Bolt) rm(c *ishell.Context) error {
	var force bool
	var args [][]byte
	for _, arg := range c.Args {
		if arg == "-f" || arg == "--force" {
			force = true
		} else {
			args = append(args, toB(arg))
		}
	}

	if len(args) == 0 {
		return fmt.Errorf("usage: rm or del [-f|--force] <key1> ... <keyN>")
	}

	return s.bolt.Update(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		for _, arg := range args {
			b := tx.Bucket(arg)
			if b != nil {
				c.Println("delete bucket:", toS(arg))
				if force {
					err = tx.DeleteBucket(arg)
				} else {
					err = forceRequired(arg)
				}
			} else {
				c.Println("delete key:", toS(arg))
				err = tx.Delete(arg)
			}
			if err != nil {
				break
			}
		}
		return
	})
}

func (s *Bolt) addPwd(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name: "pwd",
		Help: "display current bucket path",
		Func: s.pwd,
	})
}

func (s *Bolt) pwd(c *ishell.Context) {
	c.Println(sprintPath(s.current))
}

func (s *Bolt) addSequence(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:    "seq",
		Aliases: []string{"sequence"},
		Help:    "display the current sequence of bucket",
		Func:    wrapCmd(s.sequence),
	})
}

func (s *Bolt) sequence(c *ishell.Context) error {
	return s.bolt.View(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		c.Println(tx.Sequence())
		return nil
	})
}

func (s *Bolt) addSetSeq(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:    "set-seq",
		Aliases: []string{"set-sequence"},
		Help:    "set the bucket sequence",
		Func:    wrapCmd(s.setSequence),
	})
}

func (s *Bolt) setSequence(c *ishell.Context) error {
	if len(c.Args) != 1 {
		return fmt.Errorf("set-seq <new sequence int>")
	}

	seq, err := strconv.ParseUint(c.Args[0], 10, 64)
	if err != nil {
		return fmt.Errorf("set-seq <new sequence int>: %v", err)
	}

	return s.bolt.Update(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		return tx.SetSequence(seq)
	})
}

func (s *Bolt) addIncr(shell *ishell.Shell) {
	shell.AddCmd(&ishell.Cmd{
		Name:    "next-seq",
		Aliases: []string{"incr", "next-sequence"},
		Help:    "increment a sequence of bucket, has alias `incr`, `next-sequence`",
		Func:    wrapCmd(s.incr),
	})
}

func (s *Bolt) incr(c *ishell.Context) error {
	return s.bolt.Update(func(tx bolt.Tx) (err error) {
		if tx, err = s.openCurrent(tx); err != nil {
			return
		}
		var id uint64
		if id, err = tx.NextSequence(); err != nil {
			return
		}
		c.Println(id)
		return
	})
}

func (s *Bolt) getKeyCompleter(adds []string, listOptions ...bolt.ListOption) func(args []string) []string {
	return func(args []string) (ret []string) {
		set, _ := s.lists(listOptions...)
		ret = adds
		for _, e := range set {
			ret = append(ret, toS(e.Key))
		}
		return
	}
}

func (s *Bolt) lists(listOptions ...bolt.ListOption) (set bolt.EntrySet, err error) {
	err = s.bolt.View(func(tx bolt.Tx) (err error) {
		if tx, err = openBucket(tx, s.current); err != nil {
			return
		}
		set, err = tx.List(listOptions...)
		return
	})
	return
}

func openBucket(tx bolt.Tx, path [][]byte) (bolt.Tx, error) {
	if len(path) > 0 {
		for _, p := range path {
			if tx = tx.Bucket(p); tx == nil {
				break
			}
		}
		if tx == nil {
			return nil, notFound(path)
		}
	}
	return tx, nil
}

func wrapCmd(run func(c *ishell.Context) error, printOK ...bool) func(c *ishell.Context) {
	return func(c *ishell.Context) {
		defer c.Println()
		defer func() {
			if re := recover(); re != nil {
				c.Println("Error:", fmt.Errorf("%v", re))
			}
		}()
		if err := run(c); err != nil {
			c.Println("Error:", err)
		} else if len(printOK) > 0 && printOK[0] {
			c.Println("OK")
		}
	}
}

func joinPath(base [][]byte, add string) [][]byte {
	switch add {
	case ".", "":
		return base
	case "..":
		if len(base) > 0 {
			return append([][]byte(nil), base[:len(base)-1]...)
		}
		return nil
	case "/":
		return nil
	default:
		return append(append([][]byte(nil), base...), toB(add))
	}
}

func sprintPath(path [][]byte) string {
	return "/" + toS(bytes.Join(path, toB("/")))
}

func notFound(path [][]byte) error {
	return fmt.Errorf("bucket not found: `%s`", sprintPath(path))
}

func forceRequired(key []byte) error {
	return fmt.Errorf("bucket delete required force arg: use `rm -f %s` to do this", string(key))
}

func isASCII(s string) bool {
	for _, c := range s {
		if c > unicode.MaxASCII {
			return false
		}
	}
	return true
}

func toS(b []byte) string {
	s := string(b)
	if isASCII(s) {
		return s
	}
	return fmt.Sprintf("%+v", b)
}

func toB(s string) []byte {
	return []byte(s)
}

func printStruct(v interface{}, lv ...int) string {
	rv, ok := v.(reflect.Value)
	if ok {
		v = rv.Interface()
	} else {
		rv = reflect.Indirect(reflect.ValueOf(v))
	}

	if s, ok := v.(fmt.Stringer); ok {
		return s.String()
	}

	if rv.Kind() != reflect.Struct {
		return fmt.Sprintf("%+v", v)
	}

	var buf bytes.Buffer
	rt := rv.Type()
	var keys []string
	var values = make(map[string]reflect.Value, rt.NumField())
	var max = 0
	for i := 0; i < rt.NumField(); i++ {
		n := rt.Field(i).Name
		keys = append(keys, n)
		values[n] = reflect.Indirect(rv.Field(i))
		if max < len(n) {
			max = len(n)
		}
	}

	var l int
	if len(lv) > 0 {
		l = lv[0]
	}
	prefix := strings.Repeat("  ", l)
	if l > 0 {
		buf.WriteByte('\n')
	}
	sMax := strconv.Itoa(max + 2)
	for _, key := range keys {
		fmt.Fprintf(&buf, "%s%-"+sMax+"s %v\n", prefix, key+":", printStruct(values[key], l+1))
	}

	return buf.String()
}
